# import pymetis

import random
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import scipy.sparse as sp
import torch
import torch.backends.mps
import torch.nn as nn
import tqdm

from data_loader import GraphDataset
import lzyutil
import LPSI
import timing


# def cluster_test():
#     graph = GraphDataset('karate')
#     rst = pymetis.part_graph(5, graph.get_adjacency())
#     print(rst[1], type(rst))


def simulate(total, ratio, cnt):
    ub = int(total * ratio)
    src_list = random.sample(range(ub), cnt)
    print(src_list)
    max_src = max(src_list)
    p = cnt / max_src
    r = 1
    this_F1 = 2 * p * r / (p + r) if (p + r) != 0 else -1
    return this_F1


def sp_edge_test():
    graph = GraphDataset('karate')
    print(graph.get_adjacency())


class SelfAttentionPooling(nn.Module):
    def __init__(self, input_dim):
        super(SelfAttentionPooling, self).__init__()
        self.attention = nn.Linear(input_dim, 1)

    def forward(self, node_features):
        attention_scores = self.attention(node_features).squeeze(-1)
        attention_scores = torch.softmax(attention_scores, dim=0)
        snapshot_embedding = torch.sum(
            attention_scores.unsqueeze(-1) * node_features, dim=0)
        return snapshot_embedding


def test_lpsi():
    data = GraphDataset('karate')
    np.set_printoptions()
    adj = data.get_adjacency_csr()[:5, :5].tocoo()
    adj = lzyutil.get_norm_laplacian_sparse(adj, True)

    print(np.array2string(adj.toarray(), threshold=np.inf, max_line_width=np.inf))

    row = [0, 0, 1, 1, 2, 2, 3, 3]
    col = [0, 1, 1, 2, 2, 3, 3, 4]
    labels = sp.coo_matrix(
        ([1] * len(row), (row, col)), shape=(4, 5)
    )
    print(LPSI.LPSI_batch_iter_process(adj, 0.3, labels, True).shape)
    print(LPSI.LPSI_batch_cached_process(adj, labels).shape)


def calc_task(arg):
    x, y = arg
    while True:
        z = x @ y
        x = z @ x
        y = z @ y
        z = x @ y
        return z


@timing.func_timer
def test_threading():
    size = 1000
    tasks = 1000
    x = np.random.random((size, size))
    y = np.random.random((size, size))

    input = [(x, y)] * tasks

    with ThreadPoolExecutor() as executor:
        executor.map(calc_task, input)


def test_tqdm():
    size = 1000
    tasks = 1000
    x = np.random.random((size, size))
    y = np.random.random((size, size))
    input = [(x, y)] * tasks
    rst = np.zeros_like(x)
    with ThreadPoolExecutor() as executor:
        for v in tqdm.tqdm(executor.map(calc_task, input), total=tasks):
            rst += v


if __name__ == '__main__':
    # test_threading()

    # 创建一个形状为 (3, 3, 3) 的示例张量
    x = torch.ones((3,4))
    y = torch.ones((3,4))
    print(torch.cat((x,y),dim=-1).shape)
